"""
将 MIDOG 2021 中  tiff 转换成 ResNet 系列需要的数据集
1: mitosis
2: not-mitosis hard negative

https://midog.deepmicroscopy.org/download-dataset/

m21
[1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149]
[2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]
[3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

f1:
    train:
        [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150
        ,3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

    test:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149]

f2:
    train:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149
        ,3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

    test:
        [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]

f3:
    train:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149
        ,2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]

    test:
        [3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

"""
import json
import os
from PIL import Image
from concurrent.futures import ThreadPoolExecutor

wsi_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/tiff_sn/"
midog_json = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/MIDOG.json"
data = json.load(open(midog_json))
images_dict = {image["file_name"]: image["id"] for image in data["images"]}
p_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/cls_128/val/p/"
n_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/cls_128/val/n/"
os.makedirs(p_path, exist_ok=True)
os.makedirs(n_path, exist_ok=True)

patch_size = 128  # 训练 image 大小
patch_offset = 128 / 2
filter_arr = [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149]

def do_convert(wsi_file):
    key = wsi_file.split("_")[0] + ".tiff"
    image_id = images_dict[key]
    if image_id not in filter_arr:
        return
    file_name = wsi_file.split('.')[0]
    file_path = os.path.join(wsi_path, wsi_file)
    img = Image.open(file_path)
    annotations = [anno for anno in data["annotations"] if
                   anno["image_id"] == image_id]
    if len(annotations) > 0:
        for anno in annotations:
            bbox_arr = anno["bbox"]
            h = int((bbox_arr[0] + bbox_arr[2]) / 2)
            w = int((bbox_arr[1] + bbox_arr[3]) / 2)
            patch = img.crop([h - patch_offset, w - patch_offset, h + patch_offset, w + patch_offset])
            category_id = anno["category_id"]
            if category_id == 1:
                patch.save(f'{p_path}/{file_name}_{w}_{h}.png')
            else:
                patch.save(f'{n_path}/{file_name}_{w}_{h}.png')

    print(f"{wsi_file} done")


with ThreadPoolExecutor(max_workers=20) as executor:
    for wsi_file in os.listdir(wsi_path):
        executor.submit(do_convert, wsi_file)
